package com.ljh.myaiagent.agent.model;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import opennlp.tools.util.StringUtil;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * 抽象基础代理类，用于管理代理状态和执行流程
 * 提供状态转换、内存管理和基于步骤的执行循环功能。
 * 子类必须实现step方法
 */

@Slf4j
@Data
public abstract class BaseAgent {
    // 核心属性
    private String name;

    // 提示
    private String systemPrompt;
    private String nextStepPrompt;

    // 状态
    private AgentState state = AgentState.IDLE;

    // 执行控制
    private int maxSteps = 10;
    private int currentStep = 0;

    // LLM 由调用方传入具体的大模型对象
    private ChatClient chatClient;

    // Memory(需自主维护会话上下文)
    private List<Message> messageList = new ArrayList<>();

    /**
     * 运行代理
     */
    public String run(String userPrompt){
        if(this.state != AgentState.IDLE){
            throw new RuntimeException("Cannot run agent from state: " + this.state);
        }
        if(StringUtil.isEmpty(userPrompt)){
            throw new RuntimeException("User prompt cannot be empty");
        }
        // 更改状态
        this.state = AgentState.RUNNING;
        // 记录上下文消息
        messageList.add(new UserMessage(userPrompt));
        // 保存结果列表
        List<String> results = new ArrayList<>();
        try {
            for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("Executing step " + stepNumber + "/" + maxSteps);
                // 单步执行
                String stepResult = step();
                String result = "Step " + stepNumber + ": " + stepResult;
                results.add(result);
            }
            // 检查是否超出步骤限制
            if(currentStep >= maxSteps){
                state = AgentState.FINISHED;
                results.add("Terminated: Reached max steps (" + maxSteps + ")");
            }
            return String.join("\n", results);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("Error executing agent", e);
            return "Error: " + e.getMessage();
        } finally {
            // 清理资源
            this.cleanup();
        }
    }

    /**
     * 流式输出
     * @param userPrompt
     * @return
     */
    public SseEmitter runStream(String userPrompt){
        SseEmitter sseEmitter = new SseEmitter(300000L); // 设置超时时间为5分钟

        // 使用线程异步处理，避免阻塞
        CompletableFuture.runAsync(() -> {
            try {
                if(this.state != AgentState.IDLE){
                    throw new RuntimeException("Cannot run agent from state: " + this.state);
                }
                if(StringUtil.isEmpty(userPrompt)){
                    throw new RuntimeException("User prompt cannot be empty");
                }
                // 更改状态
                this.state = AgentState.RUNNING;
                // 记录上下文消息
                messageList.add(new UserMessage(userPrompt));
                try {
                    for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                        int stepNumber = i + 1;
                        currentStep = stepNumber;
                        log.info("Executing step " + stepNumber + "/" + maxSteps);
                        // 单步执行
                        String stepResult = step();
                        String result = "Step " + stepNumber + ": " + stepResult;
                        // 发送每一步的结果
                        sseEmitter.send(result);
                    }
                    // 检查是否超出步骤限制
                    if(currentStep >= maxSteps){
                        state = AgentState.FINISHED;
                        sseEmitter.send("Terminated: Reached max steps (" + maxSteps + ")");
                    }
                    sseEmitter.complete();
                } catch (Exception e) {
                    state = AgentState.ERROR;
                    log.error("Error executing agent", e);
                    try {
                        sseEmitter.send("Error: " + e.getMessage());
                        sseEmitter.complete();
                    } catch (Exception ex) {
                        sseEmitter.completeWithError(ex);
                    }
                } finally {
                    // 清理资源
                    this.cleanup();
                }
            } catch (Exception e) {
                sseEmitter.completeWithError(e);
            }
        });
        // 设置超时和完成回调
        sseEmitter.onTimeout(() -> {
            this.state = AgentState.ERROR;
            this.cleanup();
            log.warn("Agent execution timed out");
        });

        sseEmitter.onCompletion(() -> {
            if(this.state == AgentState.RUNNING){
                this.state = AgentState.FINISHED;
            }
            this.cleanup();
            log.info("Agent execution completed with state: " + this.state);
        });
        return sseEmitter;
    }

    /**
     * 执行单步步骤
     * @return
     */
    public abstract String step();

    /**
     * 清理资源
     */
    protected void cleanup(){
        // 子类可以重写该方法来清理资源
    };

}
